package com.ruoyi.draw.util;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.draw.domain.ModelConfig;
import com.ruoyi.draw.domain.PromptConfig;
import com.ruoyi.draw.domain.bo.NovelCreate;
import com.ruoyi.draw.domain.bo.RoleCreate;
import com.ruoyi.draw.domain.bo.StoryBoardCreate;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

/**
 * LLM工具类 - 通用OpenAI格式模型调用
 * 支持：SiliconFlow、Qwen、Moonshot、OpenAI等
 *
 * @author Locke
 * @date 2025-10-25
 */
@Component
public class LLMuse {

    private final RestTemplate restTemplate = new RestTemplate();

    /**
     * 优化封面提示词
     */
    public String optimizeCoverPrompt(ModelConfig modelConfig, PromptConfig promptConfig, NovelCreate novelCreate) {
        String systemPrompt = promptConfig.getCoverLlm();
        String userPrompt = buildPrompt(novelCreate);
        return callLLMAPI(modelConfig, systemPrompt, userPrompt);
    }

    /**
     * 优化角色提示词
     */
    public String optimizeRolePrompt(ModelConfig modelConfig, PromptConfig promptConfig, RoleCreate novelCreate) {
        String systemPrompt = promptConfig.getRoleLlm();
        String userPrompt = buildRolePrompt(novelCreate);
        return callLLMAPI(modelConfig, systemPrompt, userPrompt);
    }

    /**
     * 优化分镜提示词
     */
    public String optimizeStoryBoardPrompt(ModelConfig modelConfig, PromptConfig promptConfig, StoryBoardCreate storyBoardCreate) {
        String systemPrompt = promptConfig.getStoryboardLlm();
        String userPrompt = buildStoryBoardPrompt(storyBoardCreate);
        return callLLMAPI(modelConfig, systemPrompt, userPrompt);
    }

    public String optimizeStoryImgPrompt(ModelConfig modelConfig, PromptConfig promptConfig, StoryBoardCreate storyBoardCreate) {
        String systemPrompt = promptConfig.getStoryboardImg();
        String userPrompt = buildStoryBoardPrompt(storyBoardCreate);
        return callLLMAPI(modelConfig, systemPrompt, userPrompt);
    }
    /**
     * 构造基础用户提示词
     */
    private String buildPrompt(NovelCreate n) {
        return "小说名称：" + n.getName() + "\n"
                + "小说类型：" + n.getType() + "\n"
                + "小说描述：" + n.getDescription() + "\n"
                + "世界观：" + n.getWorldview() + "\n"
                + "标签：" + n.getTag() + "\n"
                + "基础提示词：" + n.getPrompt();
    }

    private String buildRolePrompt(RoleCreate n) {
        return "角色名称：" + n.getRole_name() + "\n"
                + "角色描述：" + n.getRole_detail() + "\n"
                + "角色 dress：" + n.getRole_dress() + "\n"
                + "角色 prompt：" + n.getRole_prompt() + "\n"
                + "角色 img：" + n.getRole_img() + "\n"
                + "角色 ref_img：" + n.getRef_img() + "\n";
    }

    private String buildStoryBoardPrompt(StoryBoardCreate n) {
        return n.getContent();
    }

    /**
     * 通用LLM调用实现
     */
    /**
     * 通用LLM调用实现 - 修正版
     */
    private String callLLMAPI(ModelConfig modelConfig, String systemPrompt, String userPrompt) {
        String apiUrl = modelConfig.getBaseUrl();
        String apiKey = modelConfig.getApiKey();
        String model = modelConfig.getModel();

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        if (apiKey != null && !apiKey.isEmpty()) {
            headers.setBearerAuth(apiKey);
        }

        // === 构造 body ===
        JSONObject body = new JSONObject();
        body.put("model", model);

        // ✅ 改为标准 message 对象数组
        JSONArray messages = new JSONArray();
        JSONObject sys = new JSONObject();
        sys.put("role", "system");
        sys.put("content", systemPrompt);

        JSONObject user = new JSONObject();
        user.put("role", "user");
        user.put("content", userPrompt);

        messages.add(sys);
        messages.add(user);
        body.put("messages", messages);

        // 可选参数（防止某些服务端要求）
        body.put("stream", false);
        body.put("temperature", 0.7);
        body.put("top_p", 0.9);
        body.put("n", 1);
        body.put("max_tokens", 4096);

        HttpEntity<String> entity = new HttpEntity<>(body.toJSONString(), headers);

        try {
            ResponseEntity<String> response = restTemplate.postForEntity(apiUrl, entity, String.class);
            if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                JSONObject res = JSONObject.parseObject(response.getBody());
                JSONArray choices = res.getJSONArray("choices");
                if (choices != null && !choices.isEmpty()) {
                    JSONObject choice = choices.getJSONObject(0);
                    if (choice.containsKey("message")) {
                        JSONObject msg = choice.getJSONObject("message");
                        return msg.getString("content");
                    } else if (choice.containsKey("delta")) {
                        // 有些流式接口返回 delta
                        JSONObject delta = choice.getJSONObject("delta");
                        return delta.getString("content");
                    }
                }
                return "调用LLM成功但未返回内容：" + response.getBody();
            }
            return "调用LLM失败：" + response.getBody();
        } catch (Exception e) {
            e.printStackTrace();
            return "LLM调用异常：" + e.getMessage();
        }
    }

}
